import torch
from torch import nn
import torch.distributed as dist
from transformers import Qwen3Config

from nanovllm.layers.activation import SiluAndMul
from nanovllm.layers.attention import Attention
from nanovllm.layers.layernorm import RMSNorm
from nanovllm.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear
from nanovllm.layers.rotary_embedding import get_rope
from nanovllm.layers.embed_head import VocabParallelEmbedding, ParallelLMHead


class Qwen3Attention(nn.Module):
    """
    Qwen3 的注意力层基于 多头注意力（Multi-Head Attention），并针对性能做了多重优化，是模型 “理解上下文” 的核心。
    底层组件：Qwen3Attention（注意力层）
    Q/K 合并投影：QKVParallelLinear 减少线性层数量，降低内存访问开销（尤其对大模型显著）。
    Q/K 单独 RMSNorm：相比传统的输入归一化，对 Q/K 单独归一化能更精准控制注意力得分的分布，提升模型稳定性。
    KV 头数缩减：num_kv_heads < num_heads（如 8 个 KV 头对应 32 个 Q 头），通过 “KV 共享” 减少计算量，是 Qwen3 高效推理的核心优化之一。
    """
    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position: int = 4096 * 32,
        head_dim: int | None = None,
        rms_norm_eps: float = 1e-06,
        qkv_bias: bool = False,
        rope_theta: float = 10000,
        rope_scaling: tuple | None = None,
    ) -> None:
        """初始化逻辑：张量并行拆分与参数配置"""
        super().__init__()
        tp_size = dist.get_world_size()
        # 1. 注意力头的张量并行拆分（总头数必须能被 TP 进程数整除）
        self.total_num_heads = num_heads  # 全局总注意力头数（所有 GPU 合计）
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size  # 当前 GPU 负责的头数
        self.total_num_kv_heads = num_kv_heads # 全局总 KV 头数（Qwen3 支持 KV 头数 < Q 头数，优化计算）
        assert self.total_num_kv_heads % tp_size == 0
        self.num_kv_heads = self.total_num_kv_heads // tp_size # 当前 GPU 负责的 KV
        # 2. 头维度计算（默认：隐藏层维度 / 全局总头数）
        self.head_dim = head_dim or hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim  # 当前 GPU 的 Q 输出维度
        self.kv_size = self.num_kv_heads * self.head_dim # 当前 GPU 的 K/V 输出维度
        self.scaling = self.head_dim**-0.5 # 注意力得分缩放因子（避免数值过大）
        # 3. 并行投影层（核心优化：Q/K/V 合并为一个线性层，减少计算量）
        self.qkv_proj = QKVParallelLinear( # 输入→Q/K/V 的合并并行投影
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=qkv_bias,
        )
        self.o_proj = RowParallelLinear( # 注意力输出→隐藏层维度的并行投影
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
        )
        self.rotary_emb = get_rope(# 4. 旋转位置编码（RoPE）：注入位置信息，支持长上下文
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=rope_theta,
            rope_scaling=rope_scaling,
        )
        self.attn = Attention( # 5. 注意力计算核心（封装多头注意力逻辑，支持 KV 头数 < Q 头数）
            self.num_heads,
            self.head_dim,
            self.scaling,
            self.num_kv_heads,
        )
        # 6. Q/K 单独归一化（Qwen3 特性：提升注意力稳定性）
        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        """前向传播（forward）：注意力计算流程"""
        # 步骤1：Q/K/V 合并投影（1次线性层替代3次，减少计算开销）
        qkv = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        # 步骤2：Q/K 按头归一化（Qwen3 关键优化：提升注意力质量）
        # Q 归一化：先reshape到[..., num_heads, head_dim]，归一化后恢复形状
        q_by_head = q.view(-1, self.num_heads, self.head_dim)
        q_by_head = self.q_norm(q_by_head)
        q = q_by_head.view(q.shape)
        # K 归一化：同 Q 的逻辑
        k_by_head = k.view(-1, self.num_kv_heads, self.head_dim)
        k_by_head = self.k_norm(k_by_head)
        k = k_by_head.view(k.shape)
        # 步骤3：注入位置编码（RoPE）：根据 positions 对 Q/K 应用旋转
        q, k = self.rotary_emb(positions, q, k)
        # 步骤4：多头注意力计算（封装在 Attention 类中，支持 KV 稀疏、因果掩码等）
        o = self.attn(q, k, v)
        # 步骤5：输出投影（并行层，将注意力输出映射回隐藏层维度）
        output = self.o_proj(o)
        return output


class Qwen3MLP(nn.Module):
    """
    底层组件：Qwen3MLP（多层感知机）
    负责模型的 “非线性特征转换”，是解码器层的另一核心组件，采用 Qwen3 专属的优化结构。
    合并 Gate-Up 投影：MergedColumnParallelLinear 将两个线性层合并，减少 GPU 内存带宽占用（大模型推理中带宽是关键瓶颈）。
    SiluAndMul 激活：相比传统的 “SiLU + 单独乘法”，合并为一个操作减少计算延迟，符合 Qwen3 的高效设计理念。
    """
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
    ) -> None:
        """初始化逻辑：合并投影与并行设计"""
        super().__init__()
        # 核心优化：将 gate_proj 和 up_proj 合并为一个并行线性层（减少计算步骤）
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size,
            [intermediate_size] * 2, # 输出两个分支：gate 和 up，维度均为 intermediate_size
            bias=False,
        )
        # 输出投影：将 intermediate_size 映射回 hidden_size（并行层）
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
        )
        # Qwen3 固定使用 SiLU 激活（SiluAndMul 即 SiLU(x) = x * sigmoid(x)）
        assert hidden_act == "silu"
        self.act_fn = SiluAndMul()

    def forward(self, x):
        """前向传播：激活与投影流程"""
        # 步骤1：合并投影→gate 和 up 分支（1次线性层输出两个张量）
        gate_up = self.gate_up_proj(x) # 形状：[batch*seq_len, 2*intermediate_size]
        # 步骤2：SiLU 激活（gate 分支激活后与 up 分支相乘）
        x = self.act_fn(gate_up) # 内部逻辑：gate = silu(gate_up[0]), up = gate_up[1], x = gate * up
        # 步骤3：输出投影→映射回隐藏层维度
        x = self.down_proj(x)
        return x


class Qwen3DecoderLayer(nn.Module):
    """
    中层组件：Qwen3DecoderLayer（解码器层）
    Qwen3 的基本 “计算单元”，串联注意力层、MLP 层与归一化层，是模型深度的体现（如 Qwen3-7B 含 32 个该层）。
    预归一化（Pre-Norm）：归一化在注意力 / MLP 之前，提升模型训练稳定性（避免梯度消失）。
    残差优化：residual 参数作为累积变量，避免每次层计算时重复创建残差张量，降低内存开销（对大模型推理至关重要）。
    """
    def __init__(
        self,
        config: Qwen3Config,
    ) -> None:
        """初始化逻辑：层内组件组装"""
        super().__init__()
        self.self_attn = Qwen3Attention(
            hidden_size=config.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            max_position=config.max_position_embeddings,
            rms_norm_eps=config.rms_norm_eps,
            qkv_bias=getattr(config, 'attention_bias', False),
            head_dim=getattr(config, 'head_dim', None),
            rope_theta=getattr(config, "rope_theta", 1000000),
            rope_scaling=getattr(config, "rope_scaling", None),
        )
        self.mlp = Qwen3MLP(
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
        )
        # 归一化层：输入归一化（注意力前）、注意力后归一化（MLP 前）
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """前向传播：残差连接与层间数据流"""
        # 步骤1：注意力层前的归一化 + 残差连接（预归一化结构）
        if residual is None:
            # 首次进入层：residual 初始为 None，直接用 hidden_states 作为残差
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states) # 输入归一化
        else:
            # 后续层/推理优化：归一化层直接返回（hidden_states, residual），减少内存拷贝
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        # 步骤2：注意力层计算 + 残差累积（residual 暂未更新，后续归一化时合并）
        hidden_states = self.self_attn(positions, hidden_states)
        # 步骤3：注意力后归一化 + 残差合并（将注意力输出合并到残差）
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        # 步骤4：MLP 层计算（输出将作为下一层的输入）
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual # 返回 MLP 输出和累积的残差


class Qwen3Model(nn.Module):
    """
    顶层组件：Qwen3Model（基础模型）
    封装 “词嵌入→多解码器层→全局归一化” 的完整流程，是 Qwen3 的 “特征提取器”。
    """
    def __init__(
        self,
        config: Qwen3Config,
    ) -> None:
        """初始化逻辑：模型骨架组装"""
        super().__init__()
        # 张量并行词嵌入层：将 token ID 映射为 hidden_size 维度的向量（多 GPU 拆分词表）
        self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size)
        # 解码器层列表：根据 config.num_hidden_layers 堆叠 N 个 Qwen3DecoderLayer
        self.layers = nn.ModuleList([Qwen3DecoderLayer(config) for _ in range(config.num_hidden_layers)])
        # 全局归一化层：所有解码器层输出后的最终归一化
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        """前向传播：全流程特征转换"""
        # 步骤1：词嵌入→将 token ID 转为向量（并行计算，多 GPU 拆分词表）
        hidden_states = self.embed_tokens(input_ids)
        # 步骤2：遍历所有解码器层→逐层提取特征
        residual = None # 初始化残差（累积所有层的输出）
        for layer in self.layers:
            hidden_states, residual = layer(positions, hidden_states, residual)
        # 步骤3：全局归一化→合并最终残差，输出特征向量
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states # 输出形状：[batch*seq_len, hidden_size]


class Qwen3ForCausalLM(nn.Module):
    """
    顶层模型：Qwen3ForCausalLM（因果语言模型）
    在 Qwen3Model 基础上添加 “语言模型头（LM Head）”，支持文本生成（自回归因果预测），是最终对外提供推理的模型。
    权重映射表（packed_modules_mapping）：解决 “模型存储时合并投影层，加载时需拆分权重” 的问题，与前文 load_model 函数的 “打包模块处理” 完全适配，确保权重加载正确。
    权重绑定（tie_word_embeddings）：LM Head 与词嵌入层共享权重，参数数量减少一半（如 7B 模型从 14B 降至 7B），同时提升生成文本的一致性。
    """
    # 关键：权重打包映射表（与前文 load_model 函数联动）
    # 作用：加载权重时，将“拆分的权重名”映射到“合并的模块”（如 q_proj → qkv_proj 的 q 分支）
    packed_modules_mapping = {
        "q_proj": ("qkv_proj", "q"),
        "k_proj": ("qkv_proj", "k"),
        "v_proj": ("qkv_proj", "v"),
        "gate_proj": ("gate_up_proj", 0),
        "up_proj": ("gate_up_proj", 1),
    }

    def __init__(
        self,
        config: Qwen3Config
    ) -> None:
        """初始化逻辑：模型与头组装"""
        super().__init__()
        self.model = Qwen3Model(config) # 基础特征提取模型
        # 并行语言模型头：将 hidden_size 映射到词表大小（多 GPU 拆分词表）
        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
        # 绑定词嵌入与 LM Head 权重（减少参数数量，避免过拟合）
        if config.tie_word_embeddings:
            self.lm_head.weight.data = self.model.embed_tokens.weight.data

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        """核心方法：推理 仅调用基础模型提取特征（不计算 Logits），适配推理流程中的 “特征缓存” 需求。"""
        # 仅提取特征，不计算 Logits（供推理时中间特征缓存）
        hidden_states = self.model(input_ids, positions)
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        """核心方法: Logits 计算 将基础模型的特征向量转为 “词表概率分布（Logits）”，用于文本生成采样。"""
        # 特征→Logits：形状 [batch*seq_len, vocab_size]，用于后续采样（如 Sampler 类）
        logits = self.lm_head(hidden_states)
        return logits
